import numpy as np
import pylab as pl
import cv2.cv as cv # openCV
import numpy.random as nr
import os
import scipy
import math
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d as ip
from matplotlib.pyplot import savefig
import csv
from time import  time as tm
import time
import Image
import snake_energy
import copy

###### Function which converts cartesian coordinates into polar coordinates
def toPolar(x,y):
    r = math.sqrt(pow(x,2) + pow(y,2))
    if x == 0:
        x = 0.0001
    if(y < 0 and x < 0):
        theta = np.arctan2(abs(y),abs(x)) + np.pi
    elif(y < 0 and x > 0):
        theta = np.arctan2(abs(y),-x) + np.pi
    else:
        theta = np.arctan2(y,x)
    return (r,theta)


####### Function which converts polar coordinates into cartesian coordinates
def toCartesian(r, theta,(Mx,My)):
    x = (int)(r * np.cos(theta) + Mx)
    y = (int)(r * np.sin(theta) + My)
    return (x,y)

###### Function for initialization of the snake points
###### M = center; r = radius; N = number of points
def initPoints((Mx,My), r, N, snake):   
    for n in range(N):
        theta = n * (2*np.pi/N)
        snake.append(toCartesian(r,theta,(Mx,My)))
    return snake


def sort(polarSnake):
    sortPS = []
    preSort = []
    thetaVal = []
    for j in range(len(polarSnake)):
        thetaVal.append(polarSnake[j][1])
        
    ###### Rotation of Snake --> start an Minimum
    x = thetaVal.index(min(thetaVal))
    for i in (range(x,len(polarSnake))):
        preSort.append(polarSnake[i])
    for i in range(x):
        preSort.append(polarSnake[i])   

    ###### Twists in the Snake must be removed
    for i in range(len(preSort)):
        if i > 0:
            if(preSort[i-1][1] < preSort[i][1]) and (preSort[i][1] > sortPS[len(sortPS)-1][1]) and preSort[i][1]-sortPS[len(sortPS)-1][1] >= 0.01 and preSort[i][1]-sortPS[len(sortPS)-1][1] <= 1.5:   
                sortPS.append(preSort[i]) 
        else:
            sortPS.append(preSort[i])
    return sortPS

###### Function for interpolation between the snake points
def interpolate(snake,j):   
    splines = []
    Mx = 0
    My = 0
    for (x,y) in snake:
        Mx = Mx + x
        My = My + y
    Mx = Mx/len(snake)
    My = My/len(snake)
    
    shiftSnake = []
    for (x,y) in snake:
        shiftSnake.append((x-Mx,y-My))
    
    polarSnake = []
    for j in range(len(shiftSnake)):
        polarSnake.append(toPolar(shiftSnake[j][0],shiftSnake[j][1]))
    polarSnake = sort(polarSnake)
    polarSnake.append((polarSnake[0][0],(2*np.pi+polarSnake[0][1])))  

    rVal = []
    thetaVal = []
    for j in range(len(polarSnake)):
        rVal.append(polarSnake[j][0])
        thetaVal.append(polarSnake[j][1])

    f = ip(thetaVal, rVal, kind = 'cubic',bounds_error=False)    
#    xnew = np.linspace(0,7, 1000)
#    plt.plot(thetaVal,rVal,'o-',xnew,f(xnew),'-')
#    plt.legend(['data', 'cubic'], loc='best')
#    plt.show()
#    plt.savefig('plot.png')

    polarSpline = []

    for i in range(len(thetaVal)-1):
        a = np.int(round((np.rad2deg(thetaVal[i+1])-np.rad2deg(thetaVal[i]))))
        if a == 0:
            b = 0
        else:
            b = ((thetaVal[i+1]-thetaVal[i])/(a)) 

        newTheta = 0
        for j in range(a-1):
            newTheta = thetaVal[i] + (j*b)
            if newTheta <= thetaVal[0]:
                newTheta += 2*np.pi
                newR = (float)(f(newTheta))
                polarSpline.append((newR, newTheta-2*np.pi))
            else:
                newR = (float)(f(newTheta))
                polarSpline.append((newR, newTheta))
        
    for (r,theta) in polarSpline:
        splines.append(toCartesian(r, theta,(Mx,My)))
    return splines

####### Load TIFF stack
image = Image.open('helamtgfp007.tif')
index = 0
image_stack = []
try:
    while 1:
        image_stack.append(image.copy())#[index] = np.array(im)
        image.seek(image.tell()+1)
        index += 1
except EOFError:
    pass # end of sequence
 

######## general things
snake=[]

for j in range(20):
    image = image_stack[j]
    image.save("image.png")
    cv_image = cv.LoadImage("image.png")
    cv_dst = cv.CreateImage(cv.GetSize(cv_image), cv.IPL_DEPTH_8U, 1)
    cv.CvtColor(cv_image, cv_dst, cv.CV_RGB2GRAY)

    (w,h)=image.size
#    j += 1
#    print index
    if j == 0:
        ####### Place initial points
        initPoints((470,400),20,15,snake) 
    
    ####### Energy functionals, usually denoted alpha, beta and gamma.
    a=[1]
    b=[0.7]
    c=[0.7]    
    
    ####### Mask for Snake
    mask=cv.CloneImage(cv_dst)
    
    ####### Loop the contour advacement
    energy = np.zeros(100)
    
    i=0
    while i<10:
        shake = []
        frame = cv.CloneImage(cv_dst)
        for (x,y) in snake:
            #######  add noice to the particle positions to escape local minima.
            shake.append(tuple((x,y)+nr.randint(-2,3,size=2))) 
            cv.Circle(frame, (x,y),1,255,cv.CV_FILLED,1)
                        
        #######  Interpolation between snake points  
        splines = interpolate(snake,i)
        for (x,y) in splines:
            cv.Circle(frame, (x,y),1,200,cv.CV_FILLED,1)    
                
        snake = shake
    
        ####### Snake energy
        E = snake_energy.snake_energy(snake,cv_dst,(a[0],b[0],c[0]))
#        print E.E
        energy[i] = copy.copy(E.E)                     
        
    ####### Contour advancement
        snake = cv.SnakeImage(mask,snake,a,b,c,(9,9),(cv.CV_TERMCRIT_ITER,10, 0.01))

        i += 1
    ####### end
    pl.plot(energy)
    pl.savefig('energy%s'%j+'.png')
    pl.close()

    ####### Save the images
    if j < 10:
        cv.SaveImage('frame00%s'%j+'.png',frame)
    elif j < 100:
        cv.SaveImage('frame0%s'%j+'.png',frame)
    else:
        cv.SaveImage('frame%s'%j+'.png',frame)
    ####### create video from pictures
os.system('mencoder mf://*frame*.png -mf w=200:h=200:fps=4:type=png -ovc lavc \-lavcopts vcodec=mpeg4:mbd=2:trell -oac copy -o output.avi')
